package org.korosoft.rusalad.syncnrun.command;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.korosoft.rusalad.syncnrun.messaging.Message;
import org.korosoft.rusalad.syncnrun.messaging.MessageRoute;
import org.korosoft.rusalad.syncnrun.messaging.Messenger;

import java.io.*;
import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Syncnrun "execute command" RPC command.
 *
 * @author Dmitry Korotkov
 * @since 1.0
 */
public class ExecCmd implements MessageDrivenCommand<ExecCmd.ExecAnswer> {
    private static final Log log = LogFactory.getLog(ExecCmd.class);
    private static final String LINE_SEPARATOR = System.getProperty("line.separator");

    private final String[] cmdArgs;
    private final String[] envp;
    private final File dir;
    private final byte[] input;
    private final UUID routeName;
    private final boolean detach;

    public ExecCmd(String[] cmdArgs, String[] envp, File dir, byte[] input, boolean detach) {
        routeName = UUID.randomUUID();
        this.cmdArgs = cmdArgs;
        this.envp = envp;
        this.dir = dir;
        this.input = input;
        this.detach = detach;
    }

    @Override
    public Collection<UUID> getRouteNames() {
        return Collections.singletonList(routeName);
    }

    @Override
    public void processMessageAtClient(MessageRoute route, Message message) {
        if (message instanceof StdOutErrMessage) {
            StdOutErrMessage m = (StdOutErrMessage) message;
            if (m.out != null) {
                System.out.print(m.out);
            }
            if (m.err != null) {
                System.err.print(m.err);
            }
            return;
        }
        throw new RuntimeException("Unexpected message: " + message.getClass().getName());
    }

    @Override
    public ExecCmd.ExecAnswer execute(final Messenger messenger) throws IOException {
        final String[] env;
        if (envp != null) {
            Map<String, String> sysEnv = new HashMap<String, String>(System.getenv());
            for (String e : envp) {
                int idx = e.indexOf('=');
                if (idx < 0) {
                    sysEnv.remove(e);
                } else {
                    sysEnv.put(e.substring(0, idx), e.substring(idx + 1));
                }
            }
            env = new String[sysEnv.size()];
            int idx = 0;
            for (Map.Entry<String, String> entry : sysEnv.entrySet()) {
                env[idx++] = String.format("%s=%s", entry.getKey(), entry.getValue());
            }
        } else {
            env = envp;
        }
        final Process p = Runtime.getRuntime().exec(cmdArgs, env, dir);
        if (detach) {
            return new ExecAnswer(0, 0);
        }
        final StringBuilder outBuilder = new StringBuilder();
        final StringBuilder errBuilder = new StringBuilder();
        final OutputStream stdIn = p.getOutputStream();
        final BufferedReader stdOut = new BufferedReader(new InputStreamReader(p.getInputStream()));
        final BufferedReader stdErr = new BufferedReader(new InputStreamReader(p.getErrorStream()));
        final AtomicInteger numberOfMessages = new AtomicInteger(0);
        final CountDownLatch messengerShutdown = new CountDownLatch(2);
        final Thread stdOutThread = new Thread(new Runnable() {
            @Override
            public void run() {
                try {
                    for (String l = stdOut.readLine(); l != null; l = stdOut.readLine()) {
                        //noinspection StringConcatenationInsideStringBufferAppend
                        synchronized (outBuilder) {
                            outBuilder.append(l).append(LINE_SEPARATOR);
                        }
                    }
                } catch (IOException e) {
                    if (log.isTraceEnabled()) {
                        log.trace(e.getMessage(), e);
                    } else {
                        log.debug(e.getMessage());
                    }
                } finally {
                    messengerShutdown.countDown();
                    try {
                        stdOut.close();
                    } catch (IOException ignored) {
                    }
                }
            }
        }, Thread.currentThread().getName() + "-out");
        final Thread stdErrThread = new Thread(new Runnable() {
            @Override
            public void run() {
                try {
                    for (String l = stdErr.readLine(); l != null; l = stdErr.readLine()) {
                        //noinspection StringConcatenationInsideStringBufferAppend
                        synchronized (errBuilder) {
                            errBuilder.append(l).append(LINE_SEPARATOR);
                        }
                    }
                } catch (IOException e) {
                    if (log.isTraceEnabled()) {
                        log.trace(e.getMessage(), e);
                    } else {
                        log.debug(e.getMessage());
                    }
                } finally {
                    messengerShutdown.countDown();
                    try {
                        stdErr.close();
                    } catch (IOException ignored) {
                    }
                }
            }
        }, Thread.currentThread().getName() + "-err");
        final Thread messengerThread = new Thread(new Runnable() {
            @Override
            public void run() {
                try {
                    while (!messengerShutdown.await(100, TimeUnit.MILLISECONDS)) {
                        processBuffers();
                    }
                    processBuffers();
                } catch (InterruptedException ignored) {
                }
            }

            private void processBuffers() throws InterruptedException {
                String o = null, e = null;
                synchronized (outBuilder) {
                    if (outBuilder.length() > 0) {
                        o = outBuilder.toString();
                        outBuilder.setLength(0);
                    }
                }
                synchronized (errBuilder) {
                    if (errBuilder.length() > 0) {
                        e = errBuilder.toString();
                        errBuilder.setLength(0);
                    }
                }
                if (o != null || e != null) {
                    messenger.getRoute(routeName).put(new StdOutErrMessage(o, e));
                    numberOfMessages.incrementAndGet();
                }
            }
        }, Thread.currentThread().getName() + "-msg");
        messengerThread.start();
        stdOutThread.start();
        stdErrThread.start();
        if (input != null) {
            stdIn.write(input);
        }
        stdIn.close();
        int exitCode = 0;
        try {
            exitCode = p.waitFor();
            stdOutThread.join();
            stdErrThread.join();
            messengerThread.join();
        } catch (InterruptedException ignored) {
        }
        return new ExecAnswer(exitCode, numberOfMessages.get());
    }

    @Override
    public String toString() {
        return Arrays.toString(cmdArgs);
    }

    private class StdOutErrMessage implements Message {
        private final String out;
        private final String err;

        private StdOutErrMessage(String out, String err) {
            this.out = out;
            this.err = err;
        }

        @Override
        public String toString() {
            return String.format("Out size: %d; Err size: %d", out == null ? 0 : out.length(), err == null ? 0 : err.length());
        }
    }

    public class ExecAnswer implements Answer {
        public final int exitCode;
        public final int numberOfMessagesSent;

        public ExecAnswer(int exitCode, int numberOfMessagesSent) {
            this.exitCode = exitCode;
            this.numberOfMessagesSent = numberOfMessagesSent;
        }

        @Override
        public String toString() {
            return String.format("Exit code: %d", exitCode);
        }
    }
}
